"""
2022.12.01 Liukb
    新筛样本提高判定准确性（10%）
    设计：
        1. 根据当前输入的 vcf，或者目标样本 S
        2. 根据 S 获得关联样本 / 批次的 vcf.gz 。（有 normal.vcf.gz 用这个，没有的用单样本目录下的 WES20FJ0001524.COMBINED.SNP.INDEL.PASS.vcf.gz，再没有就单样本 normal.vcf 进行压缩，再合并）
        3. 将所有 vcf.gz 与基准 vcf.gz 合并。(获得一个样本列表 S + 关联样本集合 A)
        4. call bed -> relation  （根据 A 过滤掉非本批关注的关系）
        合并命令： bcftools merge --force-samples -o tmp/all.vcf.gz A B
    修复重测样本一致，可能没有提示的 bug。
    

2022.10.26 Liukb
    测试数据仅通知 kb 

2022.09.13 Liukb
    优化通知样式

2022.09.06 Liukb
    1. 核查重做的样本
    2. 核查与其他批次有亲缘关系的样本
    3. 提示漏检关系
    4. 提示多余的关系


"""



import os
import sys
import argparse
import requests
import mysql.connector
import pandas as pd
from glob import glob
import subprocess as sp
from ruamel import yaml
from pysam import VariantFile
from collections import defaultdict
sys.path.insert(0,'/share/pipeline/Tools/')
import send_message as sendMsg


# 加载固定配置
with open('/share/pipeline/NGS/Assist/cfg_auto_NGS.yaml') as f:
    cfg = yaml.safe_load(f)


class Sample:
    """
        从各数据库获取样本文件等信息
    """
    def __init__(self, id):
        self.id = id
        self.sample_info = self._sample_search(id)

    @staticmethod
    def _sample_search(id) -> dict:
        """ 通过样本查询接口, 获取样本文件路径 """
        url = cfg['sample_index_url']
        req_json = {'sample': id, 'new': '0'}
        res = requests.post(url, json=req_json)
        if res.status_code != 200:
            # print(f'索引库获取样本失败: {id}')
            return {}
        return res.json()['message'][id]

    @classmethod
    def get_name(cls, id):
        """ 获取样本的名字，1. 先从业务大厅上拿。 找不到再从旧数据库拿 """
        if '-' in id:
            id = id.split('-')[-1]
        ywdt_name = cls._get_sample_name(id)
        if ywdt_name:
            return ywdt_name
        lims_name = cls._get_lims_name(id)
        if lims_name:
            return lims_name
        print(f'{id} 样本名获取失败!')
        return id

    @staticmethod
    def _get_sample_name(id: str) -> str:
        """ 从业务大厅获取样本姓名 """
        res = requests.post(f"{cfg['ywdt']['API']['main']}{cfg['ywdt']['API']['getSampleData']}", 
            headers = {'AUTHORIZATION-API' : cfg['ywdt']['password']}, json = {'barcode' : id})
        if res.status_code == 200:
            return res.json()['data']['name']
        # print(f'{id} 业务大厅获取姓名失败: {res.status_code}')
        return ''

    @staticmethod
    def _get_lims_name(id: str) -> str:
        """ 从 lims 上获取样本姓名 """
        conn = mysql.connector.connect(user=cfg['sql_db']["user"],
                                    host=cfg['sql_db']["host"],
                                    database=cfg['sql_db']["database"],
                                    port=cfg['sql_db']["port"],
                                    password=cfg['sql_db']["password"])
        cursor = conn.cursor()
        cursor.execute(f"select patient_name from {cfg['sql_db']['info']} where patient_id = '{id}'")
        data = cursor.fetchall()   # [(name,)]
        cursor.close()
        conn.close()
        if data:
            return list(data[0])[0]
        # print(f'{id} 从 lims 获取姓名失败!')
        return ''

    @staticmethod
    def get_relation(id: str) -> list:
        """ 从业务大厅获得样本关联的亲属样本 """
        res = requests.post(f"{cfg['ywdt']['API']['main']}{cfg['ywdt']['API']['getSampleRelation']}", 
            headers = {'AUTHORIZATION-API' : cfg['ywdt']['password']}, json = {'barcode' : id})
        if res.status_code == 200:
            return res.json()['data']
        else:
            print(f'{id} 亲属关系获取失败 {res.status_code}')
            return []


class king:
    """ 运行 king 软件, 分析输入的 bed 文件 """
    def __init__(self,
        king: str='king',
        plink_beds: list=[]
    ):
        self.king = king
        self.plink_beds = plink_beds

    def run(self, out_prefix='Related'):
        """ 运行 king 软件 """
        if os.path.isfile(f'{out_prefix}.kin0'):
            os.remove(f'{out_prefix}.kin0')
        print('\n', self.plink_beds, sep='\n')
        in_bed = ','.join(self.plink_beds)
        cmd = f'{self.king} -b {in_bed} --related --degree 3 --prefix {out_prefix}'
        self.cmd_run(cmd)

    @staticmethod
    def cmd_run(cmd:str):
        """ 运行命令 """
        print("="*70, cmd, "="*70, sep="\n")
        sp.run(cmd, shell=True, stdout=sp.PIPE, stderr=sp.PIPE)


class relation_base:
    """ 整理亲缘关系对数据 """
    def __init__(self, rel_paired: dict=defaultdict(str)):
        self.rel_paired = rel_paired
    
    def parse(self, rel_dict: dict):
        """ 解析关系字典中一级亲缘关系对(多次解析会覆盖)
        rel_dict:  是一个 id 与 关系的键值对，如 {'sample1': '先证者', 'sample2': '父亲'}

        return 一个字典，如 {'sample1 <=> sample2': '亲子'}
        """
        # 所有的亲缘关系对 (注意此处同胞中包含双胞胎关系对) 
        rel = {
            '先证者 <=> 先证者': '同胞',   # 一个家系订单下不同编号的先证者，视为同胞
            '先证者 <=> 父亲': '亲子',
            '先证者 <=> 母亲': '亲子',
            '先证者 <=> 哥哥': '同胞',
            '先证者 <=> 弟弟': '同胞',
            '先证者 <=> 姐姐': '同胞',
            '先证者 <=> 妹妹': '同胞',
            '先证者 <=> 儿子': '亲子',
            '先证者 <=> 女儿': '亲子',
            '弟弟 <=> 父亲': '亲子',
            '弟弟 <=> 母亲': '亲子',
            '弟弟 <=> 弟弟': '同胞',
            '弟弟 <=> 哥哥': '同胞',
            '弟弟 <=> 妹妹': '同胞',
            '弟弟 <=> 姐姐': '同胞',
            '哥哥 <=> 父亲': '亲子',
            '哥哥 <=> 母亲': '亲子',
            '哥哥 <=> 哥哥': '同胞',
            '哥哥 <=> 妹妹': '同胞',
            '哥哥 <=> 姐姐': '同胞',
            '妹妹 <=> 父亲': '亲子',
            '妹妹 <=> 母亲': '亲子',
            '妹妹 <=> 妹妹': '同胞',
            '姐姐 <=> 父亲': '亲子',
            '姐姐 <=> 母亲': '亲子',
            '姐姐 <=> 姐姐': '同胞',
            '姐姐 <=> 妹妹': '同胞',
            '儿子 <=> 儿子': '同胞',
            '女儿 <=> 女儿': '同胞',
            '儿子 <=> 女儿': '同胞',
            '丈夫 <=> 儿子': '亲子',
            '丈夫 <=> 女儿': '亲子',
            '妻子 <=> 儿子': '亲子',
            '妻子 <=> 女儿': '亲子',
        }
        for k1, r1 in rel_dict.items():
            for k2, r2 in rel_dict.items():
                if k1 == k2 or f'{k2} <=> {k1}' in self.rel_paired:
                    continue
                # 穷举关系对遍历，不用区分先后顺序
                if f'{r1} <=> {r2}' in rel:
                    self.rel_paired[f'{k1} <=> {k2}'] = rel[f'{r1} <=> {r2}']
    
    def num_rels(self) -> int:
        """ 计算当前关系对数 """
        return len(self.rel_paired)

    def stat(self):
        """ 打印各种关系对的数量 """
        stat_d = defaultdict(list)
        for k, v in self.rel_paired.items():
            stat_d[v].append(k)
        for k, v in stat_d.items():
            print(f'{k} 数量 {len(v)} : {v}')

    def pop(self, k1: str, k2: str) -> str:
        """ 判定输入的亲缘关系对是否存在 """
        if f'{k1} <=> {k2}' in self.rel_paired:
            rel = self.rel_paired[f'{k1} <=> {k2}']
            del self.rel_paired[f'{k1} <=> {k2}']
            return rel
        elif f'{k2} <=> {k1}' in self.rel_paired:
            rel = self.rel_paired[f'{k2} <=> {k1}']
            del self.rel_paired[f'{k2} <=> {k1}']
            return rel
        else:
            return ''

    def to_str(self):
        """ 关系对转格式化后的字符串 """
        res = ''
        for k, v in self.rel_paired.items():
            k1, k2 = k.split(' <=> ')
            k1 = Sample.get_name(k1)
            k2 = Sample.get_name(k2)
            res = f'{res}{k1} <=> {k2}:{v}\n'
        return res


class extend_king(king):
    def __init__(self, 
        in_vcfs: str='', 
        in_samples: str='', 
        tag: str='',
        qc_dir:str='./', 
        tmp_dir: str='./',
        out_file: str='Out.tsv', 
        demo_vcf: str='', 
        send_msg=False
    ):
        self.in_vcfs = in_vcfs
        self.qc_dir = qc_dir
        self.tag = tag
        self.tmp = tmp_dir
        self.out_file = out_file
        self.send_msg = send_msg
        self.king = cfg['tool']['king']
        self.plink_beds = []
        self.all_samples = []
        if in_samples:
            self.all_samples = in_samples.split(',')
        self.all_vcf = [demo_vcf]
        self.all_relation = relation_base()   # 记录相关的亲缘关系对
        self.bcftools = cfg['tool']['bcftools']
        self.bgzip = cfg['tool']['bgzip']

    def main(self):
        """ 整理输入的样本、批次 """
        # 先获取 self.all_vcf self.all_samples self.all_relation
        self.collect_info()
        print("="*70)
        print('[self.all_vcf]  ', self.all_vcf)
        print('[self.all_samples]  ', self.all_samples)
        print("="*70)
        if len(self.all_vcf) > 1:
            self.cmd_run(f'{self.bcftools} merge --force-samples -o {self.tmp}/all.vcf.gz {" ".join(self.all_vcf)}')
            self.plink_beds.extend(self.call_bed(f"{self.tmp}/all.vcf.gz", f'{self.tmp}/plink'))
            self.run(f'{self.qc_dir}/King')
            res_file = f'{self.qc_dir}/King.kin0'
            # king 无结果
            if not os.path.isfile(res_file):
                info = '【无亲缘关系对】'
                if self.all_relation.rel_paired:
                    info = f'【漏检】: \n{self.all_relation.to_str()}'
            else:
                df = pd.read_table(res_file, low_memory=False)
                if df.empty:
                    info = '【无亲缘关系对】'
                    if self.all_relation.rel_paired:
                        info = f'【漏检】: \n{self.all_relation.to_str()}'
                else:
                    df.insert(4, '亲缘关系鉴定结果', df.apply(self.judge, axis=1))
                    df.loc[:, 'ID1'] = df.loc[:,'ID1'].apply(lambda x: x.split(':')[-1])
                    df.loc[:, 'ID2'] = df.loc[:,'ID2'].apply(lambda x: x.split(':')[-1])
                    df.insert(4, '实际关系', df.apply(self.real_rel, axis=1))
                    df.loc[:, 'FID1'] = df.loc[:, 'ID1'].apply(lambda s: Sample.get_name(s))
                    df.loc[:, 'FID2'] = df.loc[:, 'ID2'].apply(lambda s: Sample.get_name(s))
                    # 减少干扰项 （非本批内的关系对忽略）
                    df = df.query('Kinship > 0.13 and (ID1 in @self.all_samples or ID2 in @self.all_samples)').copy()
                    df.to_csv(self.out_file, index=False, sep='\t', encoding='utf-8')
                    print(f'\nking 鉴定关系结果保存在: {self.out_file}')
                    # 未知的也当作未鉴定出来
                    info = self.check_rel(df.query("亲缘关系鉴定结果 != '未知'"))
                    # 结果为空:
                    if not info:
                        info = '【无亲缘关系对】'
                        if self.all_relation.rel_paired:
                            info = f'【漏检】: \n{self.all_relation.to_str()}'
            print(f'\n\n=== {self.tag} : \n{info}')
            if self.send_msg:
                if '_test' in self.in_vcfs:
                    userid = 'oYOKC52AaFGip9cAYbEGjZbAhnhY'
                else:
                    userid = ",".join(cfg['send_msg']['openid'])
                sendMsg.batch(title=f'{self.tag} 亲缘关系', message=info, accept=[], check=u"生物信息部", 
                    openids=userid)
        else:
            pass
            # sys.exit(f'\n没有可用的数据进行分析')

    def real_rel(self, row):
        if f'{row[1]} <=> {row[3]}' in self.all_relation.rel_paired:
            return self.all_relation.rel_paired[f'{row[1]} <=> {row[3]}']
        elif f'{row[3]} <=> {row[1]}' in self.all_relation.rel_paired:
            return self.all_relation.rel_paired[f'{row[3]} <=> {row[1]}']
        elif row[1] == row[3]:
            return '双胞胎'
        else:
            return '未知'

    def check_rel(self, df: pd.DataFrame) -> str:
        """ 核对亲缘关系结果是否与预期的一致 """
        rel_doubt = ''
        rel_more = ''
        rel_less = ''
        rel_err = ''
        res_dup_ok = ''
        res_dup = ''
        rel_ok = ''
        # 已完成核对的亲缘关系对（防止重测样本的家系关系对重复出现）
        ok_done_rel = []
        for _, row in df.iterrows():
            if [row['ID1'], row['ID2']] in ok_done_rel:
                continue
            ok_done_rel.append([row['ID1'], row['ID2']])
            ok_done_rel.append([row['ID2'], row['ID1']])
            rel = self.all_relation.pop(row['ID1'], row['ID2'])
            # 重测样本
            if row['ID1'] == row['ID2']:
                if row['亲缘关系鉴定结果'] == '双胞胎':
                    res_dup_ok = f"{res_dup_ok}{row['ID1']} {row['FID1']} 重测样本一致。\n"
                else:
                    res_dup = f"{res_dup}{row['ID1']} {row['FID1']} 重测样本, 与第一次结果不一致。\n"
                continue
            if rel == '亲子':
                if row['亲缘关系鉴定结果'] == '亲子':
                    rel_ok = f"{rel_ok}{row['FID1']} <=> {row['FID2']}:{row['亲缘关系鉴定结果']}\n"
                elif row['亲缘关系鉴定结果'] in ['同胞', '双胞胎']:
                    rel_err = f"{rel_err}{row['FID1']} <=> {row['FID2']}:{row['亲缘关系鉴定结果']}(预期关系: {rel})\n"
                else:
                    rel_less = f"{rel_less}{row['FID1']} <=> {row['FID2']}:{rel}（不符合检出阈值）\n"
            elif rel == '同胞':
                if row['亲缘关系鉴定结果'] == '亲子':
                    rel_err = f"{rel_err}{row['FID1']} <=> {row['FID2']}:{row['亲缘关系鉴定结果']}(预期关系: {rel})\n"
                elif row['亲缘关系鉴定结果'] in ['同胞', '双胞胎']:
                    rel_ok = f"{rel_ok}{row['FID1']} <=> {row['FID2']}:{row['亲缘关系鉴定结果']}\n"
                else:
                    rel_less = f"{rel_less}{row['FID1']} <=> {row['FID2']}:{rel}（不符合检出阈值）\n"
            else:
                if row['亲缘关系鉴定结果'] in ['亲子', '同胞', '双胞胎']:
                    rel_more = f"{rel_more}{row['FID1']} <=> {row['FID2']}:{row['亲缘关系鉴定结果']}\n"
                else:
                    rel_doubt = f"{rel_doubt}{row['FID1']} <=> {row['FID2']}:{row['亲缘关系鉴定结果']}\n"
        # 核查关系对象中剩余的关系对
        for k, v in self.all_relation.rel_paired.items():
            k1, k2 = k.split(' <=> ')
            k1 = Sample.get_name(k1)
            k2 = Sample.get_name(k2)
            rel_less = f"{rel_less}{k1} <=> {k2}:{v}\n"
        # 汇总
        res = ''
        if res_dup:
            res = f'{res}【重测异常】：\n{res_dup}'
        if rel_err:
            res = f'{res}【关系错误】：\n{rel_err}'
        if rel_less:
            res = f'{res}【漏检】：\n{rel_less}'
        if rel_more:
            res = f'{res}【意外发现】：\n{rel_more}'
        if rel_doubt:
            res = f'{res}【关系需人工核查】：\n{rel_doubt}'
        # res 先记录异常，没有异常时，报告正常情况
        if not res and (res_dup_ok or rel_ok):
            res = f'{res}【关系正确】：\n{res_dup_ok}{rel_ok}'
        else:
            print(f"【正确关系】:\n{res_dup_ok}{rel_ok}")
        print(f'\n{res}')
        return f'{res}'

    def check_vcf(self, normal_vcf):
        """ 根据 normal vcf 位置，解析最终可用于合并的 vcf 文件地址 """
        local_vcf = glob(f"{os.path.dirname(normal_vcf)}/*bam.raw.var.vcf.gz")
        if local_vcf:
            return local_vcf[0]
        batch_vcf = glob(f"{os.path.dirname(os.path.dirname(normal_vcf))}/normalized.vcf.gz")
        if batch_vcf:
            return batch_vcf[0]
        # 使用 normal vcf 需要拷过来再
        new_vcf = f"{self.tmp}/{os.path.basename(normal_vcf)}"
        cmd = f"""
            cp {normal_vcf} {self.tmp}/ ;
            {self.bgzip} -f {new_vcf} ;
            {self.bcftools} index {new_vcf}.gz
        """
        self.cmd_run(cmd)
        return f"{new_vcf}.gz"

    def collect_info(self):
        """ 核对输入数据涉及的样本和批次 """
        # 当前批次样本查询重测样本和已有二代结果的亲属样本
        local_samples = [*self.all_samples]
        for in_v in self.in_vcfs.split(','):
            if not os.path.isfile(in_v):
                continue
            samples = list(VariantFile(in_v).header.samples)
            local_samples.extend(samples)
            self.all_samples.extend(samples)
            if in_v not in self.all_vcf:
                self.all_vcf.append(in_v)
        for s in set(local_samples):
            s_info = Sample(s)
            # 获得需要合并的 vcf
            for _, info in s_info.sample_info.items():
                vcf = self.check_vcf(info['vcf_file'])
                if vcf not in self.all_vcf:
                    self.all_vcf.append(vcf)
            # 获得直接关系的样本
            local_rel = {}
            for r_d in Sample.get_relation(s):
                r_s_info = Sample(r_d['id'])
                # 无二代结果或者是本批内样本就跳过
                if not r_s_info.sample_info and r_d['id'] not in local_samples:
                    continue
                local_rel[r_d['id']] = r_d['relation']
                if r_d['id'] not in self.all_samples:
                    self.all_samples.append(r_d['id'])
                for _, info in r_s_info.sample_info.items():
                    vcf = self.check_vcf(info['vcf_file'])
                    if vcf not in self.all_vcf:
                        self.all_vcf.append(vcf)
            # 仅先证者或者关系不是本批次内的样本就跳过
            if len(local_rel) == 1 or all(k not in local_samples for k in local_rel):
                continue
            self.all_relation.parse(local_rel)

    def call_bed(self, vcf: str, prefix: str) -> list:
        """ 从 vcf 文件生成 plink bed """
        if not os.path.isfile(vcf):
            sys.exit(f'vcf 文件不存在: {vcf}')
        if os.path.isfile(f"{prefix}.bed"):
            os.remove(f"{prefix}.bed")
        self.cmd_run(f"{cfg['tool']['plink']} --make-bed --out {prefix} --vcf {vcf}")
        bed_file = glob(f'{prefix}*.bed')
        if bed_file:
            return(bed_file)
        return []

    @staticmethod
    def judge(Row):
        """
            使用 1587 对亲缘关系划分的阈值
        """
        if Row['HomIBS0'] > 0.23 and Row['HetConc'] < 0.38 and Row['PropIBD'] < 0.1:
            return "未知"
        if Row['Kinship'] > 0.375 and Row['HetConc'] > 0.6 and Row['HomIBS0'] < 0.1:
            return "双胞胎"
        if Row['HomIBS0'] == 0 or (Row['HetConc'] + Row['HomIBS0']) < 0.5:
            if Row['N_SNP'] < 7000 and (Row['HetConc'] + Row['HomIBS0'] + Row['PropIBD']) > 0.55:
                return "同胞"
            return "亲子"
        else:
            return "同胞"


if __name__ == '__main__':
    arg = argparse.ArgumentParser()
    arg.add_argument('--in_vcfs', default='')
    arg.add_argument('--in_samples', default='')
    arg.add_argument('--out_file', default='./OUT.tsv')
    arg.add_argument('--tag', default='TEST')
    arg.add_argument('--qc_dir', default='./')
    arg.add_argument('--tmp', default='./')
    arg.add_argument('--demo_vcf', default='/share/pipeline/NGS/qc/db/demo_vcf.vcf.gz')
    arg.add_argument('--send_msg', action = 'store_true', default=False)
    arg.print_help=lambda:print(f"""
    用法: {sys.argv[0]} [-h] [--in_vcfs IN_VCFS] [--in_samples IN_SAMPLES]
                        [--out_file OUT_FILE] [--tag BATCH]
                        [--qc_dir QC_DIR] [--tmp TMP]
                        [--demo_vcf DEMO_VCF] [--send_msg]

    逻辑：
        先将输入项所有 vcf 文件内的样本亲缘关系全部找到，将所有有亲缘关系的样本 vcf 合并为一个总的 vcf，
        利用 plink 转成 bed 格式，再使用 king 软件鉴别亲缘关系对。

    参数:
        --help           帮助信息
        --in_vcfs        多样本的 vcf 文件，多个 vcf 用英文逗号分隔 （需要压缩+索引）
        --in_samples     样本编号，多个以英文逗号分隔
        --tag          批次编号（用于打印输出）
        --out_file       结果文件, 默认为 ./OUT.tsv
        --qc_dir         关键的中间文件存放目录, 默认为 ./
        --tmp            多余的中间文件路径, 默认为 ./
        --demo_vcf       固定样本的 vcf 地址。（理论上样本量越多，关系判定越可靠）
        --send_msg       发送结果通知到微信
    """)
    opt = arg.parse_args()
    if opt.in_vcfs or opt.in_samples:
        print(opt)
        extend_king(
            opt.in_vcfs, 
            opt.in_samples, 
            opt.tag,
            opt.qc_dir, 
            opt.tmp, 
            opt.out_file, 
            opt.demo_vcf, 
            opt.send_msg
        ).main()
    else:
        print(f"\n--in_vcfs 与 --in_samples 必须要输入一个！")


